
import cv2
import pytesseract
from pytesseract import Output
import numpy as np
from openpyxl import Workbook

def enhance_table_recognition(image_path, output_excel):
    # 读取并预处理图像
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 增强对比度和降噪
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    enhanced = clahe.apply(gray)
    
    # 检测表格线
    thresh = cv2.threshold(enhanced, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3,3))
    detected_lines = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=2)
    
    # 使用表格专用OCR配置
    custom_config = r'--oem 3 --psm 6 -c tessedit_char_whitelist=0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ.,:;()%$/-'
    data = pytesseract.image_to_data(enhanced, config=custom_config, output_type=Output.DICT)
    
    # 创建Excel并填充数据
    wb = Workbook()
    ws = wb.active
    
    # 按单元格位置组织数据
    cell_data = {}
    for i in range(len(data['text'])):
        if int(data['conf'][i]) > 60:  # 只保留置信度高的结果
            x, y = data['left'][i], data['top'][i]
            cell_data[(x,y)] = data['text'][i]
    
    # 按坐标排序后写入Excel
    for idx, (pos, text) in enumerate(sorted(cell_data.items())):
        row = pos[1] // 30  # 估算行号(根据实际调整)
        col = pos[0] // 100  # 估算列号
        ws.cell(row=row+1, column=col+1, value=text)
    
    wb.save(output_excel)
    print(f"表格已保存至: {output_excel}")

# 使用示例
enhance_table_recognition('table_screenshot.png', 'output_1.xlsx')
